from typing_extensions import Self
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.models import UserMessage
from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, MultiModalMessage, StopMessage, ToolCallExecutionEvent
from autogen_core import Component, Image as AGImage
from autogen_agentchat.base import TaskResult, TerminatedException, TerminationCondition
from autogen_agentchat.conditions import ExternalTermination, MaxMessageTermination, TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_core import CancellationToken
import PIL
from autogen_core import Image
from io import BytesIO
import asyncio
import requests
from typing import Literal, Sequence
from pydantic import BaseModel


# Define a model client. You can use other model client that implements
# the `ChatCompletionClient` interface.


model_client = OpenAIChatCompletionClient(
    model="qwen-plus",  # Using qwen-plus instead of qwen-vl-max for better function calling support
    api_key="sk-925b8bbb82424b74a8de940d2dc5a6ce",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model_info={
        "vision": False,
        "function_calling": True,
        "json_output": False,
        "family": "unknown",
        "structured_output": True,
        "max_tokens": 8192,
        "top_p": 0.8
    }
)


primary_agent = AssistantAgent(
    name="primary",
    model_client=model_client,
    model_client_stream=False,
    system_message="You are a helpful assistant."
)

circle_agent = AssistantAgent(
    name="circle",
    model_client=model_client,
    model_client_stream=False,
    system_message="Provide constructive feedback for every message. Respond with 'APPROVE' to when your feedbacks are addressed"
)

async def main():

    max_msg_termination = MaxMessageTermination(max_messages=3)
    round_robin_team = RoundRobinGroupChat([primary_agent, circle_agent], termination_condition=max_msg_termination)

    # Use asyncio.run(...) if you are running this script as a standalone script.
    await Console(round_robin_team.run_stream(task="Write a unique, Haiku about the weather in Paris"))


#asyncio.run(main())

class FunctionCallTerminationConfig(BaseModel):
    """Configuration for the termination condition to allow for serialization
    and deserialization of the component.
    """

    function_name: str


class FunctionCallTermination(TerminationCondition, Component[FunctionCallTerminationConfig]):
    """Terminate the conversation if a FunctionExecutionResult with a specific name is received."""

    component_config_schema = FunctionCallTerminationConfig
    component_provider_override = "autogen_agentchat.conditions.FunctionCallTermination"
    """The schema for the component configuration."""

    def __init__(self, function_name: str) -> None:
        self._terminated = False
        self._function_name = function_name

    @property
    def terminated(self) -> bool:
        return self._terminated

    async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
        if self._terminated:
            raise TerminatedException("Termination condition has already been reached")
        for message in messages:
            if isinstance(message, ToolCallExecutionEvent):
                for execution in message.content:
                    if execution.name == self._function_name:
                        self._terminated = True
                        return StopMessage(
                            content=f"Function '{self._function_name}' was executed.",
                            source="FunctionCallTermination",
                        )
        return None

    async def reset(self) -> None:
        self._terminated = False

    def _to_config(self) -> FunctionCallTerminationConfig:
        return FunctionCallTerminationConfig(
            function_name=self._function_name,
        )

    @classmethod
    def _from_config(cls, config: FunctionCallTerminationConfig) -> Self:
        return cls(
            function_name=config.function_name,
        )

async def response_tool(approval: str = "approved") -> str:
    """Approve the message when all feedbacks have been addressed. 
    
    Args:
        approval: The approval status, default is "approved"
    
    Returns:
        A confirmation string indicating approval
    """
    print(f"called approve method with approval: {approval}")
    return f"Message {approval} successfully."


primary_agent_2 = AssistantAgent(
    "primary_2",
    model_client=model_client,
    system_message="You are a helpful AI assistant.",
)

critic_agent_2 = AssistantAgent(
    "critic_2",
    model_client=model_client,
    tools=[response_tool],  # Register the approve function as a tool.
    system_message="""You are a critic. Review messages carefully. 
    
IMPORTANT: When you approve a message, you MUST call the response_tool function using function calling (not just mention it in text).
You must use the actual function calling mechanism, not just write it as text.
After reviewing, if you approve, immediately call response_tool() using function calling format.""",
    reflect_on_tool_use=True,  # Enable reflection on tool use
)


function_call_termination = FunctionCallTermination(function_name="response_tool")
round_robin_team = RoundRobinGroupChat([primary_agent_2, critic_agent_2], termination_condition=function_call_termination)

async def test():
    # Use asyncio.run(...) if you are running this script as a standalone script.
    await Console(round_robin_team.run_stream(task="写一首关于北京秋天的诗"))
    await model_client.close()

asyncio.run(test())